import paddle.fluid as fluid
import paddle.fluid.layers as L


def randperm(n):
    if isinstance(n, fluid.Variable):
        index = L.arange(0, n, dtype=n.dtype)
        index = L.unsqueeze(index, 1)
        return fluid.contrib.layers.shuffle_batch(index)
    else:
        return L.randperm(n)